"""
@Description :   MLLM grounding 能力测试
@Author      :   tqychy 
@Time        :   2025/08/23 11:26:13
"""
import sys

sys.path.append("./testers")

import json
import os
import re

import matplotlib
import matplotlib.pyplot as plt
from base_tester import BaseTester
from matplotlib.patches import Rectangle
from tqdm import tqdm

matplotlib.use('Agg')


class RecTester(BaseTester):
    
    @staticmethod
    def add_bboxes(image, bboxes, img_save_path, colors=["red", "green"]):
        """
        可视化图像并添加 bboxes
        """
        ax = plt.gca()
        ax.imshow(image)
        for i, bbox in enumerate(bboxes):
            if bbox is None:
                continue
            color = colors[i % len(colors)]
            box_plot = Rectangle(
                (bbox[0], bbox[1]), bbox[2], bbox[3], fill=False, edgecolor=color, linewidth=3)
            ax.add_patch(box_plot)

        plt.savefig(img_save_path)
        plt.clf()

    def make_prompt(self, image_path, sentence, category):
        return self.model.make_prompt(image_path, sentence, category)

    @staticmethod
    def parse_output(output: str):
        pattern = r"\[\s*[+-]?(?:\d+\.?\d*|\.\d+)\s*,\s*[+-]?(?:\d+\.?\d*|\.\d+)\s*,\s*[+-]?(?:\d+\.?\d*|\.\d+)\s*,\s*[+-]?(?:\d+\.?\d*|\.\d+)\s*\]"
        match = re.search(pattern, output)
        return match.group() if match else None

    def test(self):
        with tqdm(total=len(self.dataset)) as pbar:
            for data in self.dataloader:
                image_path = data["image_path"][0]
                image_name = os.path.basename(image_path)
                image = self.load_image(image_path)
                gt_bbox = [float(x) for x in data["bbox"][0]]
                sentence = data["sentences"][0]
                category = data["category"]

                inputs = self.make_prompt(image_path, sentence, category)
                try:
                    outputs = self.model(inputs)
                except Exception as e:
                    self.logger.error(f"img {image_name} failed with {e}")
                    pbar.desc = f"IoU/mIoU: error/{self.metrics_handler.miou():.3f}"
                    pbar.update(1)
                    continue
                try:
                    parsed_output = self.parse_output(outputs)
                    bbox_pred = json.loads(parsed_output)
                    bbox_pred = self.covert_formatted_bbox(
                        bbox_pred, image.shape)
                except Exception as e:
                    self.logger.error(
                        f"img {image_name} failed with {e}: {outputs}")
                    bbox_pred = None

                iou, _, _, _ = self.metrics_handler.update(bbox_pred, gt_bbox)
                miou = self.metrics_handler.miou()
                ap50, ap75, ap90 = self.metrics_handler.map()
                self.logger.info(
                    f"img {image_name}: IoU {iou:.4f} mIoU {miou:.4f} AP50 {ap50:.4f} AP75 {ap75:.4f} AP90 {ap90:.4f} gt_bbox {gt_bbox} pred_bbox {bbox_pred}")
                self.add_bboxes(image, [bbox_pred, gt_bbox], os.path.join(
                    self.pred_save_path, image_name))
                pbar.desc = f"IoU/mIoU: {iou:.3f}/{self.metrics_handler.miou():.3f}"
                pbar.update(1)

        miou = self.metrics_handler.miou()
        map50, map75, map90 = self.metrics_handler.map()
        self.logger.critical(
            f"mIoU: {miou} AP50: {map50} AP75: {map75} AP90: {map90}")
